
#ifndef COMPTORQUE_H_
#define COMPTORQUE_H_


#include <rbdl/rbdl.h>
#include "robotStruct.hpp"

using namespace RigidBodyDynamics;
using namespace RigidBodyDynamics::Math;

class ComputedTorqueCont
{
public:
	ComputedTorqueCont(RobotModel* m);
	~ComputedTorqueCont();
	
	void setupBodies(RobotModel* m);

	void setKp(double Kp) { this->Kp = Kp; }
	void setKv(double Kv) { this->Kv = Kv; }




	void compute(std::vector<double> q,
					std::vector<double>& q_dot,
					std::vector<double>& q_d,
					std::vector<double>& q_dot_d,
					std::vector<double>& q_ddot_d,
					std::vector<double>& tau);


	void setNumBodies(unsigned int numBodies) { this->numBodies = numBodies; }


	void setSatUpper(double satUpper) { this->satUpper = satUpper; }
	void setSatLower(double satLower) { this->satLower = satLower; }


private:

	Model*	rModel;

	unsigned int bodyID;

	double Kp;
	double Kv;

	unsigned int numBodies;
	double tau_dash[4];			// dodgy.... should build the array



	VectorNd Q;
	VectorNd QDot;
	VectorNd Tau_dash;
	VectorNd QDDot;
	MatrixNd M;
	VectorNd N;

	VectorNd Tau_out;

	double satUpper;
	double satLower;
};

ComputedTorqueCont::ComputedTorqueCont(RobotModel* m) {

		// Rigid Body stuff
		rbdl_check_api_version (RBDL_API_VERSION);
		rModel = NULL;

		rModel = new Model();
		rModel->gravity = Vector3d (0., 0.0, -9.81);

		setupBodies(m);

}

ComputedTorqueCont::~ComputedTorqueCont() {}



// Assuming fixed size model
void ComputedTorqueCont::setupBodies(RobotModel* m) {

	Kp = 1.0;
	Kv = 1.0;
	numBodies = 4;
	satUpper = 1000.0;
	satLower = -1000.0;


	// Setup the rigid body dynamics model

	// Cylinder
	Body body_1 = Body (m->m1, Vector3d (m->c1, 0., 0.0), 
						         Matrix33d ((m->m1/12)*(3*(m->l1*m->l1) + m->h1*m->h1) , 0, 0, 						// Inertia tensor for a cylinder
						                    			0, (m->m1/12)*(3*(m->l1*m->l1) + m->h1*m->h1), 0, 
						                    			0, 0, 0.5*m->m1*m->l1*m->l1));
	Joint joint_1 = Joint(JointTypeRevolute, Vector3d (0., 1., 0.));
	unsigned int body_1_id = rModel->AddBody(0, Xtrans(Vector3d(0., 0., 0.)), joint_1, body_1);



	Body body_2 = Body (m->m2, Vector3d (m->c2, 0., 0.0), 
						         Matrix33d ((m->m2/12)*(m->h2*m->h2 * m->l2*m->l2) , 0, 0, 						// Inertia tensor for a cube thing
						                    			0, (m->m2/12)*(m->w2*m->w2 * m->h2*m->h2), 0, 
						                    			0, 0, (m->m2/12)*(m->l2*m->l2 * m->w2*m->w2)));
	Joint joint_2 = Joint(JointTypeRevolute, Vector3d (0., 0., 1.));
	unsigned int body_2_id = rModel->AddBody(body_1_id, Xtrans(Vector3d(0., m->l2, 0.)), joint_2, body_2);




	Body body_3 = Body (m->m3, Vector3d (m->c3, 0., 0.0), 
						         Matrix33d ((m->m3/12)*(m->h3*m->h3 * m->l3*m->l3) , 0, 0, 
						                    			0, (m->m3/12)*(m->w3*m->w3 * m->h3*m->h3), 0, 
						                    			0, 0, (m->m3/12)*(m->l3*m->l3 * m->w3*m->w3)));
	Joint joint_3 = Joint(JointTypeRevolute, Vector3d (0., 0., 1.));
	unsigned int body_3_id = rModel->AddBody(body_2_id, Xtrans(Vector3d(0., m->l3, 0.)), joint_3, body_3);




	Body body_4 = Body (m->m4, Vector3d (m->c4, 0., 0.0), 
						         Matrix33d ((m->m4/12)*(m->h4*m->h4 * m->l4*m->l4) , 0, 0, 
						                    			0, (m->m4/12)*(m->w4*m->w4 * m->h4*m->h4), 0, 
						                    			0, 0, (m->m4/12)*(m->l4*m->l4 * m->w4*m->w4)));
	Joint joint_4 = Joint(JointTypeRevolute, Vector3d (0., 0., 1.));
	unsigned int body_4_id = rModel->AddBody(body_3_id, Xtrans(Vector3d(0., m->l4, 0.)), joint_4, body_4);



	Q = VectorNd::Zero (rModel->dof_count);
	QDot = VectorNd::Zero (rModel->dof_count);
	Tau_dash = VectorNd::Zero (rModel->dof_count);
	QDDot = VectorNd::Zero (rModel->dof_count);

	M = MatrixNd::Zero(rModel->dof_count, rModel->dof_count);
	N = VectorNd::Zero(rModel->dof_count);

	Tau_out = VectorNd::Zero(rModel->dof_count);
}






void ComputedTorqueCont::compute(std::vector<double> q,
				std::vector<double>& q_dot,
				std::vector<double>& q_d,
				std::vector<double>& q_dot_d,
				std::vector<double>& q_ddot_d,
				std::vector<double>& tau)
{

	// Generate control signal for whole manipulator. Craig 10.4 (pg 295) 
	for(unsigned int i=0; i<numBodies; i++) {
		Tau_dash[i] = q_ddot_d[i] + Kv*(q_dot_d[i]-q_dot[i]) + Kp*(q_d[i]-q[i]);

		Q[i] = q[i];			// copy in positions and velocities
		QDot[i] = q_dot[i];
	}




	// Do computed torque control

	CompositeRigidBodyAlgorithm(*rModel, Q, M);					// Find M(q)
	std::cout << "Mass matrix: \n" << M << std::endl;
	// Page 103 - Rigid body dynamics algorithms - Featherstone
	// Given tau = M(q)qddot + N(q,qdot,fext)
	// Then, for some inverse dynamics algo:
	// tau = ID(model, q, qdot, qddot, fext)
	// Can find:
	// N = ID(model, q, qdot, 0, fext)    -(acceleration and hence M(q) = 0)
	InverseDynamics(*rModel, Q, QDot, VectorNd::Zero(rModel->dof_count), N);			// Find N from the model
	std::cout << "N2: " << N << std::endl;

	std::cout << "M*Tau_dash: " << M*Tau_dash << std::endl;

	// Normal bit - alpha = M, beta = N
	Tau_out = M*Tau_dash + N;

	std::cout << "Tau_dash: " << Tau_dash.transpose() << std::endl;
	std::cout << "From Tau_out: " << Tau_out.transpose() << std::endl;


//std::cout << "tau size " << tau.size() << ", TauOut size " << Tau_out.size() << std::endl;

	// Copy into openrave-ish vector and saturate.
	for(unsigned int i=0; i<numBodies; i++) {
		if (Tau_out[i] > satUpper) Tau_out[i] = satUpper;
		if (Tau_out[i] < satLower ) Tau_out[i] = satLower;
		tau[i] = Tau_out[i];
	}

}				 




#endif